#!/usr/bin/env python3

import argparse
import pymongo
import pandas as pd
import numpy as np
import dask.bag
import itertools
from pprint import pprint

parser = argparse.ArgumentParser(description='Calculate OMB - OMB.')
parser.add_argument('--host', help='Observation database host', required=True)
parser.add_argument('--port', help='Observation database port', required=True, type=int)
parser.add_argument('--synop', action='store_true')
parser.add_argument('--raob', action='store_true')
parser.add_argument('--amdar', action='store_true')
parser.add_argument('--profiler', action='store_true')
args = parser.parse_args()

missing_value = -1e10

def convert_missing_value(x):
	if x == missing_value:
		return np.nan
	else:
		return x

def mul(x, y):
	if x == missing_value or y == missing_value:
		return np.nan
	else:
		return x * y

if args.synop:
	def calc_synop_omb2(sid):	
		client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
		db = client.metobs
		ngb_data = {}
		try:
			for ngb in db.synop_ngb.find_one({ 'sid': sid })['ngb']:
				ngb_data[ngb['sid']] = ngb['d']
		except:
			#print(f'Cannot find ngb for {sid}')
			return []
		omb2_data = []
		ngb_sids = list(ngb_data.keys())
		print(f'==> {sid} has {len(ngb_data)} neighbors.')
		for this_omb in db.synop_era5_omb.find({ 'sid': sid }):
			for ngb_omb in db.synop_era5_omb.find({ 'sid': { '$in': ngb_sids }, 'time': this_omb['time'] }):
				omb2_data.append({
					'sid': [this_omb['sid'], ngb_omb['sid']],
					'time': this_omb['time'],
					'd_km': ngb_data[ngb_omb['sid']] / 1000,
					'omb_ua_var': mul(this_omb['omb_ua'], ngb_omb['omb_ua']),
					'omb_va_var': mul(this_omb['omb_va'], ngb_omb['omb_va']),
					'omb_ta_var': mul(this_omb['omb_ta'], ngb_omb['omb_ta']),
					'omb_sh_var': mul(this_omb['omb_sh'], ngb_omb['omb_sh']),
					'omb_p_var' : mul(this_omb['omb_p' ], ngb_omb['omb_p' ])
				})
		client.close()
		return omb2_data

	client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
	sids = [doc['sid'] for doc in client.metobs.synop_stations.find()]
	client.close()
	df = pd.DataFrame(dask.bag.from_sequence(sids[0:1], npartitions=10).map(calc_synop_omb2).compute())
	df.to_hdf('synop_omb2.h5', key='d')

if args.raob:
	def calc_raob_omb2(sid):
		client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
		db = client.metobs
		ngb_data = {}
		try:
			for ngb in db.raob_ngb.find_one({ 'sid': sid })['ngb']:
				ngb_data[ngb['sid']] = ngb['d']
		except:
			#print(f'Cannot find ngb for {sid}')
			return []
		omb2_data = []
		ngb_sids = list(ngb_data.keys())
		print(f'==> {sid} has {len(ngb_data)} neighbors.')
		for this_omb in db.raob_era5_omb.find({ 'sid': sid }):
			for ngb_omb in db.raob_era5_omb.find({ 'sid': { '$in': ngb_sids }, 'time': this_omb['time'], 'p': this_omb['p'] }):
				omb2_data.append({
					'sid': [this_omb['sid'], ngb_omb['sid']],
					'time': this_omb['time'],
					'd_km': ngb_data[ngb_omb['sid']] / 1000,
					'p': this_omb['p'],
					'omb_ua_var': mul(this_omb['omb_ua'], ngb_omb['omb_ua']),
					'omb_va_var': mul(this_omb['omb_va'], ngb_omb['omb_va']),
					'omb_ta_var': mul(this_omb['omb_ta'], ngb_omb['omb_ta']),
					'omb_sh_var': mul(this_omb['omb_sh'], ngb_omb['omb_sh'])
				})
		client.close()
		return omb2_data

	client = pymongo.MongoClient(f'mongodb://{args.host}', args.port)
	sids = [doc['sid'] for doc in client.metobs.raob_stations.find()]
	client.close()
	data = list(itertools.chain(*dask.bag.from_sequence(sids, npartitions=10).map(calc_raob_omb2).compute()))
	df = pd.DataFrame(data)
	df.to_hdf('raob_omb2.h5', key='d')

print(df)
